import numpy as np
import pylab as pl


def sine_wave(ampl=1, fc=1000, phase=0, fs=32000, duration=0.01):
    # 时间序列
    t = np.arange(duration*fs)/fs    
    # 等同于 np.arange(0, duration, 1/fs)

    wave = ampl * np.exp(1j*(2 * np.pi * fc * t + phase))
    phs_end = phase + 2 * np.pi * fc * duration  # 末相位
    return wave, phs_end


def fast_sine_est(wave, fs=32000, freq_list = [1000, 2000, 3000, 4000]): 
    '''返回正弦波的幅度、频率、和相位，和投影向量'''  
    
    # f_obs = np.arange(0, fs//2, f_res)
    f_obs = np.array(freq_list)           # 从 sine_est 到 sine_est_fast 只有这一行变化 
    N = wave.shape[0]
    T = N/fs                  # total time
    prj_coef = []
    for fr in f_obs:
        x_ref = sine_wave(ampl=1, fc=fr, phase=0, duration=T, fs=fs)[0]
        x_b = np.sum(wave * np.conj(x_ref))/N  
        prj_coef.append(x_b)
    
    f_idx = np.argmax(np.abs(prj_coef))
    F_est = f_obs[f_idx]                       # 频率    
    A_est = prj_coef[f_idx] * np.sqrt(2)       # 幅度
    Phs_est = np.angle(A_est)                  # 相位  
    return A_est, F_est, Phs_est, prj_coef


def plot_wave(x, fs=32000, display=True, hold=False, title=''):
    x = np.array(x)
    # 1 计算时间轴序列 t 
    t = np.arange(x.shape[0]) / fs

    # 2 plot
    if display:
        if not hold:
            pl.figure(figsize=(15, 2))
        if np.isreal(x[0] and x[1] and x[2]):  # 用前3个数作数实数/复数判断的标准
            pl.plot(t, x)
        else:
            pl.plot(t, x.real, label="I")
            pl.plot(t, x.imag, '--', label="Q")
            pl.legend(loc=1)

        pl.xlim([t[0], t[-1] + t[1] - t[0]])
        pl.xlabel('Time (s)')

        pl.ylabel('AMP')
        if title:
            pl.title(title)
        pl.tight_layout()
        return t, x
    else:
        return t, x  # 只要数据不画图

def rand_bits(n: int) -> str:
    """产生随机'0', '1' 代表的比特串

    为了程序简单'0', '1'只在字表示代表比特``0``, ``1``，实质是字符串.
    
    Parameters
    ----------
    n: integer
        比特数量

    Returns
    -------
    bit_out : string
        随机'1', '0'码组成的二进制数据
        
    Examples
    --------
    
    >>> rand_bits(16)
    '1111011111110110'
    
    """

    rand_num = np.random.rand(n)  # 均匀分布
    bit_out = ''
    for r in rand_num:
        if r > 0.5:
            bit_out += '1'
        else:
            bit_out += '0'
    return bit_out

def rand_wave(ampl=None, fc=None, phase=None, duration=0.01, fs=32000):
    
    # 生成随机参数
    if ampl == None:
        ampl = 0.5 + np.abs(np.random.rand())  # 有效值
    if fc == None:
        fc = 10 + np.random.rand() * fs/2     # 随机频率
    if phase == None:
        phase = np.random.rand() * 2 * np.pi  # 随机相位

    t = np.arange(0, duration, 1 / fs)
    cplx_wave = ampl * np.exp(1j * (2 * np.pi * fc * t + phase)) * np.sqrt(2)

    return cplx_wave, ampl, fc, phase

def plot_spec(x: np.ndarray, fs: int, fftsize: int = None,
              oversample: int = 1, unit='dB', fftshift=True,
              xlim=None, ylim: tuple = (-60, 20), 
              linestyle='-', label='', quiet=False):
    """绘制信号的频谱， 是对fft()的包装，与自然习惯一致。
    
    - 横轴以 0 频率为中心频率；
    - 横轴坐标标注 -fs/2 ~ fs/2.

    Parameters
    ----------    
    x : array like
        输出时域信号
    Ti : scalar value
        信号采样的时间间隔
    fftsize : int (default None)
         fft运算的点数。当不给定fftsize时，fftsize在内部会确定为全长 x.shape[0]。
    unit : {'linear', 'dB'}
        纵轴度量的尺度是均匀的，还是对数(非均匀)的。
    oversample: int
        过采样DFT分析时的过采样因子, oversample = 1 时相当于没有过采样

    Returns
    -------
    spec : array like
        频谱
    f_axis : array like
        频率轴. 范围是 [-fs/2 : fs/2 ]，fs=1/Ti, 间隔是 fs/fftsize
        
    Notes
    ----- 
     
    每次取一段数据求 FFT 谱并累加平均，数据不重叠:
        
    - 当 x 长度不足 1 个 fftsize 时，x的尾部补零到长度为 fftsize;
    - 当 x 的长度长于 1 个 fftsize 时，尾部不足以构成一个 fftsize 时舍弃? 

    当 x 的长度不足以达到fftsize*oversample 时，会有 FutureWarning 的警告，
    这可能是新版本 numpy 的问题。等待 numpy 解决。 

    """

    x = np.array(x)

    if fftsize is None:
        fftsize = x.shape[0]

    spec_sum = np.zeros(int(fftsize * oversample))

    for kk in range(max(x.shape[0] // fftsize, 1)):
        # 至少会累加一次，不足一个 fftsize 会自动补零
        # 不足整数个 fftsize 的数据会舍弃
        x_section = x[kk * fftsize: (kk + 1) * fftsize]
        spec_x = np.fft.fft(x_section, fftsize * oversample) \
                 / np.sqrt(fftsize * oversample)

        spec_sum += np.abs(spec_x)

    spec = spec_sum / (kk + 1) + np.finfo(float).eps  #

    if unit == 'dB':
        spec = 20 * np.log10(spec)

    if fftshift == True:
        spec = np.fft.fftshift(spec)
        f_axis = (np.arange(0, fftsize * oversample, 1, dtype='float32')
                  - fftsize * oversample / 2) / (fftsize * oversample) * fs

    else:
        f_axis = (np.arange(0, fftsize * oversample, 1, dtype='float32')) \
                 / (fftsize * oversample) * fs

    if quiet is False:
        pl.plot(f_axis, spec, linestyle, label=label)
        if ylim is not None:
            pl.ylim(ylim)

        if xlim is None:
            pl.xlim(f_axis.min(), f_axis.max())
        else:
            pl.xlim(xlim)

        pl.xlabel('Frequency [Hz]')
        pl.ylabel('Magnitude [dB]')

    return spec, f_axis

def sine_est(wave, fs=100, f_res = 0.1): 
    '''返回正弦波的幅度、频率、和相位，和投影向量'''  
    
    f_obs = np.arange(0, fs//2, f_res)
    N = wave.shape[0]
    T = N/fs                  # total time
    prj_coef = []
    for fr in f_obs:
        x_ref = sine_wave(ampl=1, fc=fr, phase=0, duration=T, fs=fs)[0]
        x_b = np.sum(wave * np.conj(x_ref))/N  
        prj_coef.append(x_b)
    
    f_idx = np.argmax(np.abs(prj_coef))
    F_est = f_obs[f_idx]                       # 频率    
    A_est = prj_coef[f_idx] * np.sqrt(2)     # 幅度
    Phs_est = np.angle(A_est)                  # 相位  
    return A_est, F_est, Phs_est, prj_coef

def imperfect_carrier(ampl=1, freq=1000, phase=0, fs=32000, duration=0.01,
                      ampl_error=0.05, freq_error=0.01, phase_error=0.01):
    # ampl, freq, 以及 phase 噪声都是按比例引入
    t = np.arange(duration*fs)/fs

    ampl = ampl * (1+ np.random.randn() * ampl_error)
    freq = freq * (1+ np.random.randn() * freq_error)
    phase = phase * (1+ np.random.randn() * phase_error)

    carrier = ampl * np.exp(1j*(2 * np.pi * freq * t + phase))
    phs_end = phase + 2 * np.pi * freq * duration  # 末相位
    return carrier, phs_end


def delay(wave, M):
    return np.hstack((np.zeros(M), wave))


class ImperfectCarrier:
    
    def __init__(self, ampl=1, freq=1000, phase=0, fs=32000, 
                 ampl_error=0.05, freq_error=0.01, phase_error=0.01):        
        self.ampl = ampl * (1+ np.random.randn() * ampl_error)
        self.freq = freq * (1+ np.random.randn() * freq_error)
        self.phase = phase * (1+ np.random.randn() * phase_error)
        self.phase_end = self.phase
        self.fs = fs

    def next(self, duration=0.01, freq_offset=0, phase_offset=0):
        t = np.arange(duration*self.fs)/self.fs
    
        freq = self.freq + freq_offset
        phase = self.phase_end+phase_offset
        carrier = self.ampl * np.exp(1j*(2 * np.pi * freq * t + phase))
        self.phase_end = self.phase_end + 2 * np.pi * freq * duration  # 末相位
        return carrier